# emotet-fla
from contextlib import AsyncExitStack
from multiprocessing.context import assert_spawning
import string
from typing import Set
import idc
import idaapi
import ida_kernwin
import ida_funcs
import idautils

from example.ida.utils import guess_machine,expr2colorstr
from future.utils import viewitems, viewvalues

from miasm.arch.x86.regs import *
from miasm.core.bin_stream_ida import bin_stream_ida
from miasm.expression.simplifications import expr_simp
from miasm.ir.ir import IRBlock, AssignBlock
from miasm.ir.symbexec import SymbolicExecutionEngine
from miasm.analysis.data_flow import load_from_int
from miasm.analysis.data_flow import ReachingDefinitions,DiGraphDefUse
from miasm.expression.expression import ExprLoc, ExprInt, ExprOp, ExprAssign,ExprCond, ExprId, ExprMem,ExprSlice
from miasm.analysis.simplifier import IRCFGSimplifierCommon, IRCFGSimplifierSSA
from miasm.core.locationdb import LocationDB
from miasm.core.asmblock import AsmCFG,AsmBlock,AsmConstraint
from miasm.ir.translators.translator import Translator
from miasm.ir.ir import IRCFG

from draw import *
import z3
import mylib
from importlib import reload
from BaseSolver import BaseSolver
# from mylib import ActionList
# from mylib import cfg_features_check,stop_on_jmp

DEBUG=1
SIMPLIFY=1

# 向后搜索，寻找jmp指令
def get_block_father(asmcfg, blk_offset):
    blk = asmcfg.getby_offset(blk_offset)
    checklist = [blk.loc_key]
    # print('pre',asmcfg.predecessors(blk.loc_key))
    pre_list=asmcfg.predecessors(blk.loc_key)
    # 没有前驱的块，就是函数的入口块
    if len(pre_list)==0:
        print('==0',hex(blk_offset))
        return blk_offset
    pred = pre_list[0]
    while True:
        curr_bloc = asmcfg.loc_key_to_block(pred)
        if curr_bloc.lines[-1].name in ['JZ', 'JMP', 'JNZ']:
            break
        checklist.append(pred)
        pred = asmcfg.predecessors(curr_bloc.loc_key)[0]
    return asmcfg.loc_db.get_location_offset(checklist[-1])

# 得到偏移信息
def get_cfg_info(asmcfg:AsmCFG):
    preds={}
    # 遍历基本块
    for blk in asmcfg.blocks:
        offset=asmcfg.loc_db.get_location_offset(blk.loc_key)
        # 针对每个块，寻找其前驱
        preds[offset]=asmcfg.predecessors(blk.loc_key)
    # dispatcher offset
    # emotet没有预分发器
    dispatcher=sorted(preds,key=lambda key:len(preds[key]),reverse=True)[0]
    relevant_blocks=[]
    for loc in preds[dispatcher]:
        offset=asmcfg.loc_db.get_location_offset(loc)
        relevant_blocks.append(get_block_father(asmcfg,offset))
    return relevant_blocks,dispatcher


def stop_on_jmp(mdis, cur_bloc, offset_to_dis):
    # emotet定制
    instr_num=len(cur_bloc.lines)
    if instr_num>=2:
        and_status_check=cur_bloc.lines[-2].name not in ['AND']
    else:
        and_status_check=True
    jmp_instr_check = cur_bloc.lines[-1].name in ['JMP','JZ','JNZ']

    if jmp_instr_check & and_status_check:
        cur_bloc.bto.clear()
        offset_to_dis.clear()

def find_state_var_usedefs(ircfg, search_var):
    var_addrs = set()
    reachings = ReachingDefinitions(ircfg)
    # 定义使用图
    # def-use图，包含use的变量以及全部的def,由定义指向使用
    digraph = DiGraphDefUse(reachings)
    # the state var always a leaf
    #for leaf in digraph.leaves():
    for leaf in digraph.nodes():
        if leaf.var == search_var:
            for x in (digraph.reachable_parents(leaf)):
                var_addrs.add(hex(ircfg.get_block(x.label)[x.index].instr.offset))
    return var_addrs

# 找到所有的backbone块,backbone块总是以cmp或者jz指令作为块的第一条指令
def get_backbone_blks(asmcfg:AsmCFG):
    blks_list=[]
    for blk in asmcfg.blocks:
        if (blk.lines[0].name=='CMP' and blk.lines[0].args[0])==ExprId('ESI',32) \
            or blk.lines[0].name=='JZ':
            offset=asmcfg.loc_db.get_location_offset(blk.loc_key)
            # blks_list.append(ExprInt(offset,32))
            blks_list.append(offset)
    return blks_list

# 检查路径的可行性
def check_path_feasibility(conds):
    solver = z3.Solver()
    # 遍历每个cond，lval为cond_expr，右侧为cond_expr值
    for lval, rval in conds:
        # print(Translator.to_language("c").from_expr(lval))
        z3_cond = Translator.to_language("z3").from_expr(lval)
        solver.add(z3_cond == int(rval.arg))

    rslt = solver.check()

    if rslt == z3.sat:
        # print(solver.model())
        return True
    else:
        return False

# 寻找逻辑块的第一个块
def get_pred_list(asmcfg,blk_offset):
    blk=asmcfg.getby_offset(blk_offset)
    pred_list=[]
    pred_list=asmcfg.predecessors(blk.loc_key)

# 扫描程序中的常量
# 仅扫描第一块
def constant_scanner(asmcfg:AsmCFG,constant_list,start_addr):
    blk = asmcfg.getby_offset(start_addr)
    # 仅遍历了第一块
    constant_map={}
    for statement in blk.lines:
        if statement.name=='MOV':       
            if statement.args[0] in constant_list:
                constant_map[statement.args[0]]=statement.args[1]
                constant_map[ExprId(statement.args[0].name+'_init',32)]=statement.args[1]
             
    return constant_map
# 扫描每一个asmblock
def global_constant_scanner(asmcfg:AsmCFG,constant_list,constant_map):
    global_constant_map={}
    #　init
    for constant in constant_list:
        global_constant_map[constant]=[]
    # 遍历每个块
    for asmblk in asmcfg.blocks:
        for statement in asmblk.lines:
            if statement.name=='MOV':
                if statement.args[0] in constant_list:
                    if isinstance(statement.args[1],ExprInt):
                        global_constant_map[statement.args[0]].append(statement.args[1])
                    elif isinstance(statement.args[1],ExprId):
                        # 存在常量记录
                        if statement.args[1] in list(constant_map.keys()):
                            global_constant_map[statement.args[0]].append(constant_map[statement.args[1]])
                    else:
                        pass
    for constant in constant_list:
        global_constant_map[constant]=list(set(global_constant_map[constant]))
        global_constant_map[ExprId(constant.name+'_init',32)]=list(set(global_constant_map[constant]))
    return  global_constant_map

# 得到可能的状态值
def find_var_asg(ircfg, var,constant_map,type=0):
    val_list = []
    res = {}
    # 保证对于结果来说，使用next键将不会产生错误
    res['next']=None
    for lbl, irblock in viewitems(ircfg.blocks):
        # 每条单独的语句
        for assignblk in irblock:
            result = set(assignblk).intersection(var) #检测交集，没有交集表示该语句并未涉及状态变量的操作
            if not result:
                continue
            else:
                dst, src = assignblk.items()[0]
                if isinstance(src, ExprInt):
                    res['next'] = int(src)
                    val_list += [int(src)]
                elif isinstance(src, ExprSlice):
                    # 得到其前驱块列表
                    pass
                elif isinstance(src, ExprId):
                    # 使用常量扫描的结果
                    if src in constant_map.keys():
                        if type==0:
                            res['next']=int(constant_map[src])
                        elif type==1:
                            constant_list=constant_map[src]
                            res['next']=[]    
                            for n in constant_list:
                                res['next'].append(int(n))
                    else:
                        pass
                else:
                    print('unknow src type')
    return res, val_list



def get_solver_result(conds):
    solver = z3.Solver()
    for lval, rval in conds:
        # print(Translator.to_language("c").from_expr(lval))
        z3_cond = Translator.to_language("z3").from_expr(lval)
        solver.add(z3_cond == int(rval.arg))
    rslt = solver.check()
    if rslt==z3.sat:
        return solver.model()
    else:
        return None
    
# 路径探索
DEPTH_LIMIT=20
class FinalState:
    def __init__(self, result, sym, path_conds, path_history,out_addr=None):
        self.result = result
        self.sb = sym
        self.path_conds = path_conds
        self.path_history = path_history
        self.out_range_addr= out_addr
def expoler(asmcfg,ir,start_addr,start_symbols,
            ircfg,constant_map,cond_limit=7,uncond_limit=100,
            lbl_stop=None,final_states=[]):
    # print('do')
    blk_list=get_backbone_blks(asmcfg)
    def expoler_path(addr,symbols,conds,depth,final_staates,path):
        if depth>DEPTH_LIMIT:
            # print('[*]unexcepted depath,check maybe wrong',addr)
            # print(path)
            # final_states.append(FinalState(False, None, conds, path))
            return
        # 初始化符号执行引擎
        # print(hex(addr))
        sb=SymbolicExecutionEngine(ir,symbols)
        for _ in range(uncond_limit):
            if isinstance(addr, ExprInt):
                # print(type(addr),addr)
                if int(addr) not in blk_list:
                    reslt=get_solver_result(conds)
                    # print(addr,reslt)
                    final_states.append(FinalState(True, sb, conds, path, addr))
                    # print('sorry but i need stop',addr)
                    return
            # 判断该地址的最后一条语句是否为JMP
            path.append(addr)
            # print(hex(addr))
            pc=sb.run_block_at(ircfg, addr)
            
            if isinstance(pc, ExprCond):
                # print('pc',pc,pc.cond.args)
                # pc_replace=pc.replace_expr({ExprId('EBP_init',32):ExprInt(1,32)})
                pc_replace=pc.replace_expr(constant_map)
                pc=pc_replace
                # try:
                #     # print(pc.cond.op)
                #     tmp=pc.cond.args[0].args[0].args
                #     print(type(tmp),tmp[0],tmp[1])
                #     if isinstance(tmp[0],ExprId) and isinstance(tmp[1],ExprId):
                #         print(tmp)
                #         assert()
                # except Exception as e:
                #     print(e.__str__())
                #     assert()
                
                # print('pc',pc)
                cond_true={pc.cond:ExprInt(1,32)}
                cond_false={pc.cond:ExprInt(0,32)}
                # print(pc)
                addr_true=expr_simp(sb.eval_expr(pc.replace_expr(cond_true),{}))
                addr_false=expr_simp(sb.eval_expr(pc.replace_expr(cond_false),{}))
                # print('t',addr_true)
                # print('f',addr_false)
                conds_true=list(conds)+list(cond_true.items())
                conds_false=list(conds)+list(cond_false.items())
                sat_true=check_path_feasibility(conds_true)
                sat_false=check_path_feasibility(conds_false)
                
                # expoler_path(addr_true, sb.symbols.copy(), 
                #         conds_true, depth + 1, final_states, list(path))
                # expoler_path(addr_false, sb.symbols.copy(), 
                #         conds_false, depth + 1, final_states, list(path))
                if sat_true:
                    expoler_path(addr_true, sb.symbols.copy(), 
                        conds_true, depth + 1, final_states, list(path))
                else:
                    # final_states.append(FinalState(False, sb, conds_true, path))
                    pass
                if sat_false:
                    expoler_path(addr_false, sb.symbols.copy(), 
                        conds_false, depth + 1, final_states, list(path))
                else:
                    # final_states.append(FinalState(False, sb, conds_false, path))
                    pass
                return
            else:
                addr = expr_simp(sb.eval_expr(pc))
            #print('path',path)
        final_states.append(FinalState(True, sb, conds, path))
        return
    return expoler_path(start_addr,start_symbols,[],0,final_states,[])
    
# 寻找完整的逻辑块
def get_logic_blk(asmcfg:AsmCFG,blk_offset):
    
    logic_list=[]
    blk=asmcfg.getby_offset(blk_offset)
    pred_list=asmcfg.predecessors(blk.loc_key)
    # print(hex(blk_offset),len(pred_list))
    if len(pred_list)<=1:
        return None
    # 判断前驱块的最后一条命令是否都不是跳转
    # 如果都不是跳转，则逻辑块的起始应该在前驱块
    for pred_blk_locKey in pred_list:
        pred_blk_offset=asmcfg.loc_db.get_location_offset(pred_blk_locKey)
        pred_blk=asmcfg.getby_offset(pred_blk_offset)
        # 判断每个块的最后一条指令 以及倒数第二条指令是否为AND
        pred_blk_len=len(pred_blk.lines)
        if pred_blk_len>=2:
            if pred_blk.lines[-2].name in ['AND']:
                and_check=True
            else:
                and_check=False
        else:
            and_check=False
        if (pred_blk.lines[-1].name in ['JZ','JNZ','JG','JNG']) and not and_check:
            # 该块的前驱块并非逻辑块
            return None
            break
        else:
            logic_offset=asmcfg.loc_db.get_location_offset(pred_blk.loc_key)
            logic_list.append(logic_offset)
    # 此处只向前寻找一次
    # 所有的最后一条语句都不是跳转
    
    return logic_list
    
    

def deal_with_muti_statement(ircfg,state_var,ir,addr):
    
    status_name=state_var.name
    # if addr==0x10021184:
    # 遍历ircfg中的每个块
    blk_keys=list(ircfg.blocks)
    for irblk_key in reversed(blk_keys):
        irblk=ircfg.blocks[irblk_key]
        #　遍历块中的每一条语句
        for single_instr in reversed(irblk):
            # 遍历块中的每一条IR表达式
            for single_expr in reversed(viewitems(single_instr)):
                # print('single_expr',single_expr[0],single_expr[1])
                # 判断操作对象的类型
                if isinstance(single_expr[0],ExprId) and (status_name in single_expr[0].name):
                    # print(type(single_expr[1]),single_expr[1].op,single_expr[1])
                    stand_expr=single_expr[1].args[-1]
                    if not isinstance(stand_expr,ExprInt):
                        continue
                    plus_expr=single_expr[1].args[-2].args[-1]
                    if isinstance(plus_expr,ExprInt):
                        plus_value=int(plus_expr)
                        stand_value=int(stand_expr)
                        
                        return {'True':plus_value+stand_value,'False':stand_value}
    return {'next':None}
    # 倒数第二条语句
    blks=list(ircfg.blocks)
    irblock= ircfg.blocks[blks[-1]]
    status_expr=viewitems(irblock[-2])[-1] # 最后一块的倒数第二条指令的最后一条表达式
    # print(status_expr[1])
    plus_expr=status_expr[1].args[0].args[1]
    stand_expr=status_expr[1].args[1]
    if isinstance(plus_expr,ExprInt):
        plus_value=int(plus_expr)
    else:
        plus_value=0
    stand_value=int(stand_expr)
    return {'True':plus_value+stand_value,'False':stand_value}
    
    # for dst,src in viewitems(irblock[-2]):
    #     # print(dst,'<-',src)
    #     if isinstance(src,ExprOp):
    #         print('*',dst,'<-',src)
    #         print(src.args[0].args[1])
    #         try:
    #             print(src.args[1])
    #         except:
    #             pass
    #         # print(src.args[0].args[2])
    return
    # 符号执行
    symbols_init={}
    for i, r in enumerate(all_regs_ids):
        symbols_init[r] = all_regs_ids_init[i]
    final_states=[]
    _sb=SymbolicExecutionEngine(ir,symbols_init)
    for _ in range(100):
        pc=_sb.run_block_at(ircfg,addr,step=True)
            
    
    
    return

# 裁剪IRCFG图
# fixed_cfg:逻辑块的状态变量值,log_rel_map:逻辑块到相关块的映射,status_map:状态值到下一跳的映射
def cut_asmcfg(action:mylib.ActionList,asmcfg:AsmCFG,fixed_cfg,log_rel_map,status_map):
    asmcfg_res=asmcfg.copy()
    # 删除无用的块
    backbone_blks=get_backbone_blks(asmcfg)
    for backbone_offset in backbone_blks:
        backbone_blk=asmcfg.getby_offset(backbone_offset)
        asmcfg_res.del_block(backbone_blk)
        action.graph_del_blk(backbone_offset)
    #　仅操作逻辑块，而不操作相关块
    # print(fixed_cfg)
    for log_blk_off in fixed_cfg:
        status_value=fixed_cfg[log_blk_off]
        # print(log_blk,':',status_value)
        # rel_blk_addr=log_rel_map[log_blk]
        # rel_blk=asmcfg_res.getby_offset(rel_blk_addr)
        log_blk=asmcfg_res.getby_offset(log_blk_off)
        try:
            # 仅有一个方向
            # print('next',status_value)
            if status_value['next'] is not None:
                src_blk_key=log_blk.loc_key
                dst_blk_addr=status_map[status_value['next']]
                dst_blk=asmcfg_res.getby_offset(dst_blk_addr)
                dst_blk_key=dst_blk.loc_key
                asmcfg_res.add_edge(src_blk_key,dst_blk_key,AsmConstraint(dst_blk_key,'c_next'))
                action.graph_add_edge(log_blk_off,dst_blk_addr)
                # print('[+]add_edge',src_blk_key,dst_blk_key)
        except Exception as e:
            print(e.__str__())
            #raise(e)
            pass
        
        try:
            # print('True',status_value)
            if status_value['True'] is not None :
                # true dst
                src_blk_key=log_blk.loc_key
                true_blk_addr=status_map[status_value['True']]
                true_blk=asmcfg_res.getby_offset(true_blk_addr)
                true_blk_key=true_blk.loc_key
                asmcfg_res.add_edge(src_blk_key,true_blk_key,AsmConstraint(true_blk_key,'c_next'))
                action.graph_add_edge(log_blk_off,true_blk_addr)
                # print('[o]add_edge',src_blk_key,true_blk_key)
                
        except Exception as e:
            print(e.__str__())
            # raise(e)
        try:
            # false dst
            # print('False',status_value)
            if status_value['False'] is not None:
                false_blk_addr=status_map[status_value['False']]
                false_blk=asmcfg_res.getby_offset(false_blk_addr)
                false_blk_key=false_blk.loc_key
                asmcfg_res.add_edge(src_blk_key,false_blk_key,AsmConstraint(false_blk_key,'c_to'))
                action.graph_add_edge(log_blk_off,false_blk_addr)
                # print('[x]add_edge',src_blk_key,false_blk_key)
        except:
            pass

    return asmcfg_res
    # for irblk in blks:
        
    
    # for irblk in blks: 
    # for irb in list(viewvalues(ircfg.blocks)):
    #     irs=[]
    #     for assignblk in irb:
    #         new_assignblk={
    #             expr_simp(dst):expr_simp(src)
    #             for dst,src in viewitems(assignblk)
    #         }
    #         # print('assignblk:',assignblk)
    #         # print('new_assignblk:',new_assignblk)
    #         irs.append(AssignBlock(new_assignblk,instr=assignblk.instr))
    #     ircfg.blocks[irb.loc_key]=IRBlock(loc_db,irb.loc_key,irs)
    return
    
def leaves_scanner(asmcfg:AsmCFG,leaves_off_list,mdis,loc_db,ir_arch,lifter,state_var,constant_map):
    more_fixed={}
    for leaf_off in leaves_off_list:
        leaf_blk=asmcfg.getby_offset(leaf_off)
        # 判断最后一条指令
        last_instr=leaf_blk.lines[-1]
        # 该块为最后一块
        if last_instr.name in ['RET','RETN']:
            continue
        else:
            # logic_off=leaf_off
            logic_off=mylib.cfg_search_pred_logic(asmcfg,leaf_off)
            asmcfg_blk=mdis.dis_multiblock(logic_off)
            ircfg_blk=lifter.new_ircfg_from_asmcfg(asmcfg_blk)
            mark_addrs=find_state_var_usedefs(ircfg_blk,state_var)
            if len(mark_addrs)==1:
                var_asg,tmpval_list=find_var_asg(ircfg_blk,{state_var},constant_map)
                if var_asg['next']:
                    more_fixed[leaf_off]=var_asg
            elif len(mark_addrs)>1:
                head=loc_db.get_offset_location(logic_off)
                ssa_simplifier=IRCFGSimplifierSSA(lifter)
                ssa=ssa_simplifier.ircfg_to_ssa(ircfg_blk, head)
                ssa_simplifier.do_propagate_expressions(ssa, head)
                var_asg=deal_with_muti_statement(ircfg_blk,state_var,ir_arch,logic_off)
                more_fixed[leaf_off]=var_asg
            else:
                print('len == 0',hex(leaf_off),leaf_blk.loc_key)
                # assert()
            pass
        # print(hex(leaf_off),last_instr) 
    print(more_fixed)       
    return more_fixed
def usedef_scanner(sanncer_blocks,rel_blk_info,fixed_cfg,mdis,loc_db,lifter,state_var,constant_map,ir_arch):
    for blk_addr in sanncer_blocks:
        asmcfg_blk=mdis.dis_multiblock(blk_addr)
        ircfg_blk=lifter.new_ircfg_from_asmcfg(asmcfg_blk)
        # ircfg_simplifier = IRCFGSimplifierCommon(lifter)
        # ircfg_simplifier.simplify(ircfg, addr)
        # ircfg_simplifier.simplify(ircfg, blk_addr)
        
        # 标记影响状态变量的指令
        mark_addrs=find_state_var_usedefs(ircfg_blk,state_var)
        # rel_blk_info[addr]=(asmcfg_blk,mark_addrs)
        rel_blk_info[blk_addr]=(asmcfg_blk,mark_addrs)
        
        # 统计
        # count_map[len(mark_addrs)]=count_map[len(mark_addrs)]+1
        # get_logic_addr(asmcfg,blk_addr)
        
        if len(mark_addrs)==1:
            var_asg,tmpval_list=find_var_asg(ircfg_blk,{state_var},constant_map)
            if var_asg['next']:
            # 保存状态变量值
                # var_asg['next']=hex(var_asg['next'])
                # fixed_cfg[hex(blk_addr)]=var_asg
                fixed_cfg[blk_addr]=var_asg
            print(hex(blk_addr),var_asg)
        # 处理影响变量数量大于1的块
        elif len(mark_addrs)>1:
            head=loc_db.get_offset_location(blk_addr)
            ssa_simplifier=IRCFGSimplifierSSA(lifter)
            ssa=ssa_simplifier.ircfg_to_ssa(ircfg_blk, head)
            ssa_simplifier.do_propagate_expressions(ssa, head)
            # graph=GraphMiasmIR(ircfg_blk,'ircfg'+hex(blk_addr),None)
            # graph.Show()
            var_asg=deal_with_muti_statement(ircfg_blk,state_var,ir_arch,blk_addr)
            # fixed_cfg[hex(blk_addr)]=var_asg
            fixed_cfg[blk_addr]=var_asg
            print(hex(blk_addr),var_asg)
            pass
        else:
            # graph=GraphMiasmIR(ircfg_blk,'else'+hex(blk_addr),None)
            # graph.Show()
            pass
        # 简化
        if DRAW_GRAPH:
            if len(mark_addrs)>1:
                graph=GraphMiasmIR(ircfg_blk,'ircfg'+hex(blk_addr),None)
                graph.Show()
    return

# 非常依赖特征的处理，填充
def emotet_fixed(asmcfg:AsmCFG,leaves_list,global_constant_map,rel_blk_info,mdis,loc_db,lifter,state_var,ir_arch):
    more_fixed={}
    for leaf_off in leaves_list:
        leaf_blk=asmcfg.getby_offset(leaf_off)
        # 判断最后一条指令
        last_instr=leaf_blk.lines[-1]
        # 该块为最后一块
        if last_instr.name in ['RET','RETN']:
            continue
        else:
            asmcfg_blk=mdis.dis_multiblock(leaf_off)
            ircfg_blk=lifter.new_ircfg_from_asmcfg(asmcfg_blk)
            # 标记影响状态变量的指令
            mark_addrs=find_state_var_usedefs(ircfg_blk,state_var)
            rel_blk_info[leaf_off]=(asmcfg_blk,mark_addrs)
            if len(mark_addrs)==1:
                var_asg,tmpval_list=find_var_asg(ircfg_blk,{state_var},global_constant_map,type=1)
                print(var_asg)
                # 此处的next应该为数组，因为find_var_asg的type为1
                if var_asg['next']:
                    _count=len(var_asg['next'])
                    if _count==1:
                        more_fixed[leaf_off]={'next':var_asg['next'][0]}
                    elif _count==2:
                        more_fixed[leaf_off]={'True':var_asg['next'][0],'False':var_asg['next'][1]}
                    else:
                        assert()
                    # var_asg['next']=hex(var_asg['next'])
                    # fixed_cfg[hex(blk_addr)]=var_asg
        return more_fixed
            
        
    # 修复mov esi,ebx导致的无法识别的常量
    for off in leaves_list:
        logic_off=mylib.cfg_search_pred_logic(asmcfg,off)
        print('logic offset',hex(logic_off))
    # 向前寻找前驱块，从相关快找到逻辑块
    return
#--------------------
SHOW_FNAME=False
DRAW_GRAPH=False
#--------------------
# def fla_emotet():
def fla_emotet(func_now,machine,asmcfg,head,ircfg,lifter,mdis,loc_db):
    reload(mylib)
    action=mylib.ActionList()
    if DEBUG:
        # leaves_off_list=[268572753, 268573146, 268571436, 268572334, 268572621, 268572987, 268573108, 268572312, 268572770, 268572904, 268571397, 268571997, 268571926, 268572273, 268572403, 268571251, 268571901, 268571603, 268572211, 268571341, 268571516, 268571227, 268572824]
        # leaves_scanner(asmcfg,leaves_off_list)
        # assert()
        pass
    mdis.dis_block_callback=stop_on_jmp
    func_addr=func_now.start_ea
    for irb in list(viewvalues(ircfg.blocks)):
        irs=[]
        for assignblk in irb:
            new_assignblk={
                expr_simp(dst):expr_simp(src)
                for dst,src in viewitems(assignblk)
            }
            # print('assignblk:',assignblk)
            # print('new_assignblk:',new_assignblk)
            irs.append(AssignBlock(new_assignblk,instr=assignblk.instr))
        ircfg.blocks[irb.loc_key]=IRBlock(loc_db,irb.loc_key,irs)
    
    # Common简化
    title = "Miasm IR graph(sim)"
    if SIMPLIFY:
        ircfg_simplifier=IRCFGSimplifierCommon(lifter)
        ircfg_simplifier.simplify(ircfg, head)
    
    # 从asmcfg中得到:相关块，分发器的偏移
    rel_blocks,dispatcher_off=get_cfg_info(asmcfg)
    dispatcher_blk=asmcfg.getby_offset(dispatcher_off)
    dispatcher_fisrt_instr=dispatcher_blk.lines[0]
    # 状态变量
    state_var=dispatcher_fisrt_instr.get_args_expr()[0]
    if DEBUG:
        print('[*]state_var',state_var)
    
    # 主干分支的信息
    backbone={}
  
    # 状态变量的值
    val_list={}

    
    count_map=[0]*20
    # 常量
    constant_map=constant_scanner(asmcfg,[ExprId('EBP',32)],func_addr)
    global_constant_map=global_constant_scanner(asmcfg,[ExprId('EBX',32)],constant_map)
    # 传播至init类型
    
    #**************************
    # 进行路径探索
    ir_arch=machine.ira(loc_db)
    symbols_init={}
    for i,r in enumerate(all_regs_ids):
        symbols_init[r]=all_regs_ids_init[i]
    # symbols_init[ExprId('ESI',32)]=ExprInt(0x2DEE833, 32) # 该值应该作为变量
    final_states=[]     # 最终的状态
    
    # 路径探索
    status_map={}
    if True:
        path_len=[0]*100
        expoler(asmcfg,ir_arch, dispatcher_off, symbols_init, ircfg,constant_map,30,100,None,final_states)
        for statee in final_states:
            path_len[len(statee.path_conds)]+=1
            # if len(statee.path_conds)==10:
            #     print(statee.path_history)
            #     print(statee.path_conds)
        print('expoler complete all',len(final_states))
        for final_item in final_states:
            # print('final conds',final_item.path_conds)
            result=get_solver_result(final_item.path_conds)
            result_list=result.__str__()[1:-1].split(',')
            for item in result_list:
                if item[0:3]=='ESI':
                    status_value=int(item.split('=')[1].strip())
                    status_map[status_value]=int(final_item.out_range_addr)
    # print status
    for key_value in status_map:
        print(hex(key_value),'->',hex(status_map[key_value]))   
                    
    # return
    
    # 从相关块得到逻辑块
    # 此处等待优化
    logic_blocks=[]
    log_rel_map={}
    for blk_addr in rel_blocks:
        # 判断在块中，状态变量是否完整
        _asmcfg_blk=mdis.dis_multiblock(blk_addr)
        _ircfg_blk=lifter.new_ircfg_from_asmcfg(_asmcfg_blk)
        # 使用的地址
        mark_addrs=find_state_var_usedefs(_ircfg_blk,state_var)
        # 分析的结果
        var_asg,tmpval_list=find_var_asg(_ircfg_blk,{state_var},constant_map)
        if var_asg['next']==None:
            tmp_blk_lst=get_logic_blk(asmcfg,blk_addr)
            if(tmp_blk_lst):
                # 连接到新的块
                logic_blocks=logic_blocks+tmp_blk_lst
                for log_blk_off in tmp_blk_lst:
                    log_rel_map[log_blk_off]=blk_addr
            else:
                logic_blocks.append(blk_addr)
                log_rel_map[blk_addr]=blk_addr
        else:
            logic_blocks.append(blk_addr)
            log_rel_map[blk_addr]=blk_addr
    # print(log_rel_map)
    #print('rel',rel_blocks)
    #print('log',logic_blocks)
    print('rel-log',set(rel_blocks)-set(logic_blocks))
    print('log-rel',set(logic_blocks)-set(rel_blocks))
    print('relvant count',len(rel_blocks))
    print('logic count',len(logic_blocks))
    
    
    # 相关块的信息
    rel_blk_info={}
    # 修复后的asmCFG
    fixed_cfg={}
    # usedef scanner
    usedef_scanner(logic_blocks,rel_blk_info,fixed_cfg,mdis,loc_db,lifter,state_var,constant_map,ir_arch)
    
    # fixed_cfg:逻辑块的状态变量值,log_rel_map:逻辑块到相关块的映射,status_map:状态值到下一跳的映射
    asmcfg_ret_1=cut_asmcfg(action,asmcfg,fixed_cfg,log_rel_map,status_map)
    
    # the second scan
    # get leaf node 
    leaves_list=mylib.cfg_features_check(asmcfg_ret_1,'out',0)
    #fix leaves_list
    
    more_fixed=leaves_scanner(asmcfg,leaves_list,mdis,loc_db,ir_arch,lifter,state_var,constant_map)
    asmcfg_ret_2=cut_asmcfg(action,asmcfg_ret_1,more_fixed,log_rel_map,status_map)
    # assert()
    global_fixed=emotet_fixed(asmcfg_ret_2,mylib.cfg_features_check(asmcfg_ret_2,'out',0),
                 global_constant_map,rel_blk_info,mdis,loc_db,lifter,state_var,ir_arch)
    asmcfg_ret=cut_asmcfg(action,asmcfg_ret_2,global_fixed,log_rel_map,status_map)
    draw_graph(asmcfg_ret,'asmcfg')
    # print(count_map)
    return action

# fla_emotet()

class EmotetSolver(BaseSolver):
    def __init__(self) -> None:
        super().__init__('Emotet')
        pass
    
    def Solve(self,obj):
        super().Solve(obj)
        self.lift()
        self.action=fla_emotet(self.func_now,self.machine,self.asmcfg,self.head,
                   self.ircfg,self.lifter,self.mdis,self.loc_db)
        return 
if __name__=='__main__':
    reload(mylib)